'''
Created on Jul 19, 2013

@author: a.renduchintala
'''

import numpy as np
from numpy.oldnumeric.random_array import multivariate_normal
from pprint import pprint
import sys

class GMM(object):
    '''
    classdocs
    '''
    
    def __init__(self, num_of_gaussians):
        '''
        Constructor
        '''
        self.K = num_of_gaussians
        self.model_estimates = {}
        self.M = 0
        self.N = 0
        self.threshold = 1e-3
        
    def initializeParams(self, training_instances):
        for k in range(self.K):
            init_sigma = np.matrix(np.diag(np.ones(self.N)))
            self.model_estimates[k] = (float(1 / self.K), np.matrix(np.random.uniform(np.min(training_instances, 0), np.max(training_instances, 0))), init_sigma)
                
        return self.model_estimates;
    
    def fit(self, training_instances):
        '''
        intializing parameters, 
        with equal weights coefs, random means and unit variance
        '''
        [self.M, self.N] = np.shape(training_instances)
        '''
        M = number of training istances
        M = number of features per instance
        '''
        self.initializeParams(training_instances)
        ll = 0.0
        oll = 1.0
        while abs(oll - ll) > self.threshold:
            oll = ll
            contributions = self.expectationStep(training_instances)
            self.model_estimates = self.maximizationStep(contributions, training_instances)
            # self.prettyprintParams()
            ll = self.loglikelihood(training_instances)
            print 'log-likelihood', ll
            
        return self.model_estimates
      
    def prettyprintParams(self):
        for k in range(self.K):
            print '\nGMM', k
            (wt, mu, sigma) = self.model_estimates[k]
            print 'weight:', wt
            print 'mean:', pprint(mu)
            print 'sigma:', pprint(sigma)
            
    def expectationStep(self, training_instances):
        x = training_instances
        '''
        compute contribution of each gaussian to a training instance
        '''
        contributions = np.matrix(np.zeros((self.M, self.K)))
        for  m in range(self.M):
            cks = []
            for k in range(self.K):
                wt = self.model_estimates[k][0]
                mu = self.model_estimates[k][1]
                sigma = self.model_estimates[k][2]
                cks.append(self.multivariateNormalPDF(x[m, :], mu, sigma))
            cks = np.array(cks, float) / np.sum(cks)
            contributions[m, :] = np.matrix(cks)
            
        return contributions
    
    def maximizationStep(self, contributions, training_instances):
        x = training_instances
        Nks = np.sum(contributions, 0)
        wt_ks = np.array(Nks, float) / self.M
        mu_ks = contributions.T * training_instances
        for k in range(self.K):
            Nk = Nks[:, k]
            mu_ks[k, :] = (mu_ks[k, :]).astype(float) / Nk
            sigma_k = np.matrix(np.zeros((self.N, self.N)))
            for i in range(self.M):
                sigma_k = sigma_k + (contributions[i, k] * (x[i, :] - mu_ks[k, :]).T * (x[i, :] - mu_ks[k, :])) 
            sigma_k = sigma_k / Nk
            if np.linalg.det(sigma_k) < 0.25:
                print 'resetting sigma_k'
                sigma_k = np.matrix(np.identity(self.N))  * 0.5
                
            # print k , wt_ks[:, k] , mu_ks[k, :] , sigma_k
            self.model_estimates[k] = (wt_ks[:, k], mu_ks[k, :], sigma_k)   
        return self.model_estimates
    
    def loglikelihood(self, X):
        log_likelihood = 0.0
        for m in range(self.M):
            cost_k = 0.0
            for k in range(self.K):
                (wt, mu, sigma) = self.model_estimates[k]
                pk = self.multivariateNormalPDF(X[m, :], mu, sigma)
                cost_k = cost_k + (wt * pk)
            log_likelihood = log_likelihood + np.log(cost_k)
        return  log_likelihood
    
    def multivariateNormalPDF(self, x, mu, sigma):
        try:
            const = np.linalg.det(sigma) ** (-0.5) * (2 * np.pi) ** (-float(self.N) / 2)  
            pdf = const * np.exp(-0.5 * (x - mu) * np.linalg.inv(sigma) * (x - mu).T)
            return pdf.item(0)
        except:
            print x , mu , sigma
            print 'exiting...'
            sys.exit(1)
            
